# -*- coding: utf-8 -*-
"""
Created on Sun Oct 11 10:09:35 2020
用于双量子点量子模型
单量子比特的归零操作
多个初始态对多个目标态的归零保真度
选用两种策略：
先采用 要么最优，要么最差
如果效果不好就选第二种
动作要么选最优的，要么选次优的

@author: Waikikilick
"""

import numpy as np
from scipy.linalg import expm
from time import *
import multiprocessing as mp
import copy
np.random.seed(1)
T = 2*np.pi
dt = np.pi/5
step_max = T/dt
sx = np.mat([[0, 1], [1, 0]], dtype=complex) 
sz = np.mat([[1, 0], [0, -1]], dtype=complex)
action_space = np.array([0,1,2,3])#,5,6,7,8])

theta_num = 6 #除了 0 和 Pi 两个点之外，点的数量
varphi_num = 21#varphi 角度一圈上的点数
#总点数为 theta_num * varphi_num + 2(布洛赫球两极)

theta = np.linspace(0,np.pi,theta_num+2,endpoint=True) 
varphi = np.linspace(0,np.pi*2,varphi_num,endpoint=False) 

def psi_set():
    psi_set = []
    for ii in range(1,theta_num+1):
        for jj in range(varphi_num):
            psi_set.append(np.mat([[np.cos(theta[ii]/2)],[np.sin(theta[ii]/2)*(np.cos(varphi[jj])+np.sin(varphi[jj])*(0+1j))]]))
    psi_set.append(np.mat([[1], [0]], dtype=complex))
    psi_set.append(np.mat([[0], [1]], dtype=complex))
    return psi_set

target_set = psi_set()
init_set = psi_set()
#----------------------------------------------------------------------------------------------------
#动作直接选最优的
def step0(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        
        H = float(action_space[action])* sz + 1 * sx
        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        
        psi_list.append(psi_)
        fid_list.append(fid)
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_

#动作选最优的，或者最差的
def step1(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        
        H = float(action_space[action])* sz + 1 * sx
        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        
        psi_list.append(psi_)
        fid_list.append(fid)
    
    if F < max(fid_list):
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    else:
        
        best_action = fid_list.index(min(fid_list))
        best_fid = min(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_

#动作选最优的，或者次优的
def step2(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        
        H = float(action_space[action])* sz + 1 * sx
        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        
        psi_list.append(psi_)
        fid_list.append(fid)
        
    if F < max(fid_list):
        best_action = fid_list.index(max(fid_list))
        best_fid = fid_list[best_action]
        
    else:
        psi_list_ = copy.deepcopy(psi_list)
        fid_list_ = copy.deepcopy(fid_list)
        
        del psi_list_[fid_list_.index(max(fid_list_))]
        del fid_list_[fid_list_.index(max(fid_list_))]
        
        best_action = fid_list.index(max(fid_list_))
        
        best_fid = max(fid_list_)
        
    psi_ = psi_list[best_action]
    
    return best_action, best_fid, psi_
#---------------------------------------------------------------------------------
#将测试集的保真度从小到大排列出来，来展示保真度分布
def sort_fid(test_fidelity_list):
    sort_fid = []
    for i in range (test_fidelity_list.shape[0]):
        b = test_fidelity_list[i,:]
        sort_fid  = np.append(sort_fid,b)
    sort_fid.sort()
    return sort_fid
#--------------------------------------------------------------------------------


def job(target_psi):
    fids_list = []
    start_time = time()
    for psi1 in init_set:
        
        psi = psi1
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        
        fid_max = F
        fid_max1 = F
        fid_max2 = F
        fid_max0 = F
        
        step_n = 0
        while True:
            action, F, psi_ = step1(psi,target_psi,F)
            fid_max1 = max(F,fid_max1)
            psi = psi_
            step_n += 1
            if fid_max1>0.999 or step_n>step_max:
                break
            
        step_n = 0
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        psi = psi1
        while True:
            action, F, psi_ = step2(psi,target_psi,F)
            fid_max2 = max(F,fid_max2)
            psi = psi_
            step_n += 1
            if fid_max2>0.999 or step_n>step_max:
                break 
            
        step_n = 0
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        psi = psi1
        while True:
            action, F, psi_ = step0(psi,target_psi,F)
            fid_max0 = max(F,fid_max0)
            psi = psi_
            step_n += 1
            if fid_max0>0.999 or step_n>step_max:
                break 
            
        fid_max = max(fid_max1,fid_max2,fid_max0)  
        fids_list.append(fid_max)
        
    return  np.mean(fids_list)

def multicore():
    pool = mp.Pool()
    F_list = pool.map(job, target_set)
    return F_list
    

if __name__ == '__main__':
    # print(target_set)
    time1 = time()
    F_list = multicore()
    print(F_list)
    time2 = time()
    print(np.mean(F_list))
    print('time_cost is: ',time2-time1)
    
    #F_list.sort()# jiang liebiao an shunxu pailie.

#目标点集
# [matrix([[0.92387953+0.j],
#          [0.38268343+0.j]]),
#  matrix([[0.92387953+0.j        ],
#          [0.27059805+0.27059805j]]),
#  matrix([[9.23879533e-01+0.j        ],
#          [2.34326020e-17+0.38268343j]]),
#  matrix([[ 0.92387953+0.j        ],
#          [-0.27059805+0.27059805j]]),
#  matrix([[ 0.92387953+0.00000000e+00j],
#          [-0.38268343+4.68652041e-17j]]),
#  matrix([[ 0.92387953+0.j        ],
#          [-0.27059805-0.27059805j]]),
#  matrix([[ 9.23879533e-01+0.j        ],
#          [-7.02978061e-17-0.38268343j]]),
#  matrix([[0.92387953+0.j        ],
#          [0.27059805-0.27059805j]]),
#  matrix([[0.70710678+0.j],
#          [0.70710678+0.j]]),
#  matrix([[0.70710678+0.j ],
#          [0.5       +0.5j]]),
#  matrix([[7.07106781e-01+0.j        ],
#          [4.32978028e-17+0.70710678j]]),
#  matrix([[ 0.70710678+0.j ],
#          [-0.5       +0.5j]]),
#  matrix([[ 0.70710678+0.00000000e+00j],
#          [-0.70710678+8.65956056e-17j]]),
#  matrix([[ 0.70710678+0.j ],
#          [-0.5       -0.5j]]),
#  matrix([[ 7.07106781e-01+0.j        ],
#          [-1.29893408e-16-0.70710678j]]),
#  matrix([[0.70710678+0.j ],
#          [0.5       -0.5j]]),
#  matrix([[0.38268343+0.j],
#          [0.92387953+0.j]]),
#  matrix([[0.38268343+0.j        ],
#          [0.65328148+0.65328148j]]),
#  matrix([[3.82683432e-01+0.j        ],
#          [5.65713056e-17+0.92387953j]]),
#  matrix([[ 0.38268343+0.j        ],
#          [-0.65328148+0.65328148j]]),
#  matrix([[ 0.38268343+0.00000000e+00j],
#          [-0.92387953+1.13142611e-16j]]),
#  matrix([[ 0.38268343+0.j        ],
#          [-0.65328148-0.65328148j]]),
#  matrix([[ 3.82683432e-01+0.j        ],
#          [-1.69713917e-16-0.92387953j]]),
#  matrix([[0.38268343+0.j        ],
#          [0.65328148-0.65328148j]]),
#  matrix([[1.+0.j],
#          [0.+0.j]]),
#  matrix([[0.+0.j],
#          [1.+0.j]])]

# #对应的保真度
[0.9985281367379384, 0.9990919567471859, 0.9970130622032529, 0.9958261601009446, 0.9952238284078838, 0.9957531079642764, 0.9950831555787869, 0.9937914188533712, 0.9917651574449484, 0.9926446441823851, 0.9876565935823789, 0.9865262384614621, 0.9783146366483826, 0.9566958286943457, 0.9391167701117846, 0.9541436487769428, 0.9832260816801249, 0.9899602974099075, 0.9948565032237372, 0.9978374519026485, 0.998901429370501,
 0.9943256369365048, 0.9979653669858046, 0.9901920927270942, 0.9892809101862989, 0.9928472764689806, 0.9880798626634254, 0.9905812889685016, 0.9904179764917787, 0.990087974947468, 0.9880456154489774, 0.9920021313156344, 0.9932630121865904, 0.9868395876024776, 0.9680604079775859, 0.9430525313278049, 0.913955817960879, 0.8618923217467215,  0.8029979413461314, 0.8053269286429415, 0.9916361446207604, 0.995203431200752, 
 0.9970375835383355, 0.9962825763179544, 0.9961302817347395, 0.991922177570806, 0.9832519916199174, 0.9723729718572476, 0.9751981038405101, 0.9749736825024166, 0.9752860845115607, 0.975148349539015, 0.9797474645786822, 0.9946591102148334, 0.9910632663448915, 0.9868498181836134, 0.9742271459983906,  0.95252337464162, 0.9191977280296351, 0.9148350255091822, 0.9213310363363949, 0.9402266190422941, 0.9857090542253423, 
 0.9859197186188758, 0.9958746583898332, 0.9891550444574094, 0.9807275481116837, 0.9666055402996148, 0.9346810558183347, 0.9128285912634242, 0.9169013006319959, 0.9266809776411367, 0.9543692704182647, 0.9955907074419881, 0.9970928091044657, 0.9964622571371108,  0.994638644455115, 0.9887826030866224, 0.9783546636106615, 0.9712884158296031, 0.9751393366486045, 0.9769088474346244, 0.9735693073208894, 0.9730592111688536, 
 0.9937676273926024, 0.9913519965898672, 0.9831122380223434, 0.9512837026973926, 0.9290987543112664, 0.8867135590581142, 0.8322626314084723, 0.7958825741307194, 0.8659195144968583, 0.9942918752211857, 0.9958440803854672,  0.9978374716101988, 0.9948588234749781, 0.9886713083786718, 0.9917457822170557, 0.9906584010557165, 0.9879311535948092, 0.9915138878255234, 0.9901828000227642, 0.9891243414340886, 0.9903879222325604, 
 0.9879877630899399, 0.9849465441087155, 0.9695690881609693, 0.9439923616109787, 0.9478903606195141, 0.9720567656634505, 0.991519965866841, 0.9938087432541696, 0.9968625424072276,  0.9977343667024143, 0.9986924760677096, 0.9985954305589945, 0.9982528441398146, 0.9960215284563367, 0.9948666378924379, 0.9955510293506391, 0.9950244042078172, 0.9951970616355128, 0.9920927830963493, 0.9930621071793576, 0.991359857558628, 
 0.9921063624582807, 0.9916140761584241]
0.9728270302083435
time_cost is:  32.18448352813721


# # #保真度排序
F_list.sort()# 将列表按顺序排列.
print(F_list)

[0.7958825741307194, 0.8029979413461314, 0.8053269286429415, 0.8322626314084723, 0.8618923217467215, 0.8659195144968583, 0.8867135590581142, 0.9128285912634242, 0.913955817960879, 0.9148350255091822, 0.9169013006319959, 0.9191977280296351, 0.9213310363363949, 0.9266809776411367, 0.9290987543112664, 0.9346810558183347, 0.9391167701117846, 0.9402266190422941, 0.9430525313278049, 0.9439923616109787, 0.9478903606195141, 0.9512837026973926, 0.95252337464162, 0.9541436487769428, 0.9543692704182647, 0.9566958286943457, 0.9666055402996148, 0.9680604079775859, 0.9695690881609693, 0.9712884158296031, 0.9720567656634505, 0.9723729718572476, 0.9730592111688536, 0.9735693073208894, 0.9742271459983906, 0.9749736825024166, 0.9751393366486045, 0.975148349539015, 0.9751981038405101, 0.9752860845115607, 0.9769088474346244, 0.9783146366483826, 0.9783546636106615, 0.9797474645786822, 0.9807275481116837, 0.9831122380223434, 0.9832260816801249, 0.9832519916199174, 0.9849465441087155, 0.9857090542253423, 0.9859197186188758, 0.9865262384614621, 0.9868395876024776, 0.9868498181836134, 0.9876565935823789, 0.9879311535948092, 0.9879877630899399, 0.9880456154489774, 0.9880798626634254, 0.9886713083786718, 0.9887826030866224, 0.9891243414340886, 0.9891550444574094, 0.9892809101862989, 0.9899602974099075, 0.990087974947468, 0.9901828000227642, 0.9901920927270942, 0.9903879222325604, 0.9904179764917787, 0.9905812889685016, 0.9906584010557165, 0.9910632663448915, 0.9913519965898672, 0.991359857558628, 0.9915138878255234, 0.991519965866841, 0.9916140761584241, 0.9916361446207604, 0.9917457822170557, 0.9917651574449484, 0.991922177570806, 0.9920021313156344, 0.9920927830963493, 0.9921063624582807, 0.9926446441823851, 0.9928472764689806, 0.9930621071793576, 0.9932630121865904, 0.9937676273926024, 0.9937914188533712, 0.9938087432541696, 0.9942918752211857, 0.9943256369365048, 0.994638644455115, 0.9946591102148334, 0.9948565032237372, 0.9948588234749781, 0.9948666378924379, 0.9950244042078172, 0.9950831555787869, 0.9951970616355128, 0.995203431200752, 0.9952238284078838, 0.9955510293506391, 0.9955907074419881, 0.9957531079642764, 0.9958261601009446, 0.9958440803854672, 0.9958746583898332, 0.9960215284563367, 0.9961302817347395, 0.9962825763179544, 0.9964622571371108, 0.9968625424072276, 0.9970130622032529, 0.9970375835383355, 0.9970928091044657, 0.9977343667024143, 0.9978374519026485, 0.9978374716101988, 0.9979653669858046, 0.9982528441398146, 0.9985281367379384, 0.9985954305589945, 0.9986924760677096, 0.998901429370501, 0.9990919567471859]
